#include <string>
#include <vector>
#include <list>

#include <m3n/random_field.h>
#include <m3n/functional_m3n.h>
#include <m3n/regressors/regressor_includes.h>

#include <iostream>

#include "opencv2/imgproc/imgproc.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/video/video.hpp"
#include <opencv2/gpu/gpu.hpp>

using namespace std;

using namespace cv;


#include "LcBasic.h"

class LcInterface
{
public:

	LcInterface();
	void input( int argc, char * const argv[]);

	//LcClassifier*  getClassifier(){ return my_classifier;}

	//LcVideoReadExt getTrainVideo();
	//LcVideoReadExt getTestVideo();


	bool IF_TRAIN;
	bool IF_TEST;
	
	bool IF_NODE_BIAS;
	bool IF_EDGE_BIAS;
	
	double nbr_iters;
	double step_size;

	int train_f_start;
	int train_f_end;
	int train_f_rate;

	int test_f_start;
	int test_f_end;
	int test_f_rate;

	string root;

	string train_video_name;
	string test_video_name;

	string feature_code_string;

	string classifier_name;

	void buildDatasetmap();

	map< string , Vec3i> data_map;

	bool setVideo( string name ,int & start, int & end , int & rate);
	void setClassifier();

	void init();
	

	//LcClassifier* my_classifier;
	


};

//LcVideoReadExt LcInterface::getTrainVideo()
//{
//	LcVideoReadExt ans;
//	ans.video_name = train_video_name;
//	ans.f_start = train_f_start;
//	ans.f_end = train_f_end;
//	ans.f_rate = train_f_rate;
//	ans.t_win = 0;
//	ans.root = root;
//	return ans;
//}
//
//LcVideoReadExt LcInterface::getTestVideo()
//{
//	LcVideoReadExt ans;
//	ans.video_name = test_video_name;
//	ans.f_start = test_f_start;
//	ans.f_end = test_f_end;
//	ans.f_rate = test_f_rate;
//	ans.t_win = 0;
//	ans.root = root;
//	return ans;
//}

void LcInterface::setClassifier()
{
	if(classifier_name == "rdt")
	{
		//my_classifier = new LcRandomTrees;
	}
}
void LcInterface::init()
{
	if(IF_TEST == false)
	{
		if(train_f_start<0) train_f_start = test_f_start;
		if(train_f_end<0) train_f_end = test_f_end;
		if(train_f_rate<0) train_f_rate = test_f_rate;
	}

	if(feature_code_string == "") feature_code_string = "l";

	if(classifier_name == "" ) classifier_name = "rdt";
	setClassifier();

	if(train_video_name == "" ){train_video_name = "GOPR2382_cab";}
	setVideo( train_video_name, train_f_start, train_f_end, train_f_rate);

	if(test_video_name == "" ){test_video_name = "GOPR2384_cab";}
	setVideo( test_video_name, test_f_start, test_f_end, test_f_rate);
}

LcInterface::LcInterface()
{
	train_f_start = -1;
	train_f_rate = -1;
	train_f_end = -1;

	test_f_start = -1;
	test_f_rate = -1;
	test_f_end = -1;

	root = "./../_FGS/";
	train_video_name = "";
	test_video_name = "";
	feature_code_string = "";
	classifier_name = "";
	
	IF_NODE_BIAS = true;
	IF_EDGE_BIAS = true;
	nbr_iters = 50;
	step_size = 0.1;

	IF_TEST = false;
	IF_TRAIN = false;

	buildDatasetmap();
}

void LcInterface::buildDatasetmap()
{
	data_map.clear();
	data_map.insert( pair< string, Vec3i>( "GOPR2382_cab", Vec3i(150,5300,50)) );
	data_map.insert( pair< string, Vec3i>( "GOPR2384_cab", Vec3i(150,11200,25)) );
}

bool LcInterface::setVideo( string name ,int & start, int & end , int & rate)
{
	map<string , Vec3i>::iterator iter = data_map.find( name );
	if( iter != data_map.end())
	{
		Vec3i temp = iter->second;
		if(start<0)start = temp[0];
		if(end<0)end = temp[1];
		if(rate<0)rate = temp[2];
		return true;
	}
	else
	{
		cout << "One of Video is not in default" << endl;
		return false;
	}
}

void LcInterface::input( int argc, char * const argv[])
{
	for(int ind = 1;ind < argc;ind++)
	{
		/*** Parse switches ***/
		if (argv[ind][0] == '-') {
			switch (argv[ind][1]) {
				case 'f':
					root = string(argv[++ind]);
					break;
				case 't':
					train_video_name = string(argv[++ind]);
					break;
				case 'p':
					test_video_name = string(argv[++ind]);
					break;
				case 'v':
					feature_code_string =  string(argv[++ind]);
					break;
				case 'a':
					IF_TRAIN = true;
					break;
				case 'b':
					IF_TEST = true;
					break;
				case 's':
					//only for predict video setting , for trainning video it's always solid with frame
					// unless it's no testing model then the test frame setting will be use
					// as training parameter
					test_f_start =  atoi(argv[++ind]);
					break;
				case 'e':
					test_f_end = atoi(argv[++ind]);
					break;
				case 'r':
					test_f_rate = atoi(argv[++ind]);
					break;
				case 'c':
					classifier_name = string(argv[++ind]);
					break;
				case 'i':
					IF_NODE_BIAS = false;
					break;
				case 'j':
					IF_EDGE_BIAS = false;
					break;
				case 'k':
					step_size = atof(argv[++ind]);
					break;
				case 'l':
					nbr_iters = atoi(argv[++ind]);
					break;	
				default : printf("Unknown switch '%c'\n", argv[ind][1]);
				break;
			}
		}
	}

	init();
}

int main(int argc, char * const argv[])
{
	float label_threshold = 0.5f;

	LcInterface my_setting;
	
	my_setting.input( argc, argv);
	
	bool IF_NODE_BIAS = my_setting.IF_NODE_BIAS;
	
	bool IF_EDGE_BIAS = my_setting.IF_EDGE_BIAS;

	string root = my_setting.root;//"./../../_FGS/";
	string video_name = my_setting.train_video_name;//"GOPR2393_cab";
	string src_video = root + "vid/"+ video_name+".avi";
	
	int start_frame= my_setting.train_f_start;

	int end_frame = my_setting.train_f_end;

	int rate_frame = my_setting.train_f_rate;
	
	vector< Mat > node_descs;
	vector< Mat > edge_descs;
	vector< Mat > edge_maps;
	vector< Mat > labels;
	
	{
		int _n = (end_frame - start_frame) / rate_frame;
		node_descs.reserve( _n + 5 ); node_descs.clear();
		edge_descs.reserve( _n + 5 ); edge_descs.clear();
		edge_maps.reserve( _n + 5 );  edge_maps.clear();
		labels.reserve( _n + 5 );     labels.clear();
	}
	
	for(int i = start_frame;i<end_frame;i+=rate_frame)
	{
		cout << "=============" << i << "=================" << endl;
		
		{
			Mat node_desc;
			char fname_node_desc[400];
			sprintf(fname_node_desc,"%s/superpixel_output/%s_node_desc_%d.bin"
				 ,root.c_str(), video_name.c_str(),i);
			lc::LcBin2Mat( fname_node_desc, node_desc);
			
			node_descs.push_back( node_desc );
		}
		
		{
			Mat edge_desc;
			Mat edge_map;
			
			char fname_edge_desc[400];
			sprintf(fname_edge_desc,"%s/superpixel_output/%s_edge_desc_%d.bin"
				,root.c_str(), video_name.c_str(),i);
			lc::LcBin2Mat( fname_edge_desc, edge_desc);

			edge_map.convertTo(edge_map,5);
			char fname_edge_map[400];
			sprintf(fname_edge_map,"%s/superpixel_output/%s_edge_map_%d.bin"
				,root.c_str(), video_name.c_str(),i);
			lc::LcBin2Mat( fname_edge_map, edge_map);
			
			edge_descs.push_back( edge_desc );
			edge_maps.push_back( edge_map );
		}
		
		{
			Mat label;		

			char fname_label[400];
			sprintf(fname_label,"%s/superpixel_output/%s_label_%d.bin"
				,root.c_str(), video_name.c_str(),i);
			lc::LcBin2Mat( fname_label, label);
			
			labels.push_back( label );
		}
	}
	
	// count summary number of nodes;
	
	int n_node = 0;
	
	for(int i = 0 ; i < (int) node_descs.size() ; i++ )
		n_node += node_descs[i].rows;
		
	int n_edge = 0;
	
	for(int i = 0 ; i < (int) edge_descs.size(); i++)
		n_edge += edge_descs[i].rows;
		
	int dim_node = node_descs[0].cols;
	
	int dim_edge = edge_descs[0].cols;
	
	if(IF_NODE_BIAS) dim_node++;
	
	if(IF_EDGE_BIAS) dim_edge++;
	
	cout << "dim of node " << dim_node << endl;
	cout << "dim of edge " << dim_edge << endl;
	
	unsigned int nbr_clique_sets = 1; // edges, high-order cliques
	
	RandomField tmp_rf(nbr_clique_sets);
	
	vector< const RandomField::Node* > nodes;
	
	nodes.reserve( n_node );
	
	vector< vector < double > > node_features;
	vector< vector < double > > edge_features;
	
	node_features.resize( n_node );
	edge_features.resize( n_edge );
	
		
	{
		// add all node feature
		
		int node_shift = 0;	
		
		for(int i = 0 ; i < (int) node_descs.size() ; i++ )
		{
			int _rows = node_descs[i].rows;
			int _cols = node_descs[i].cols;
					
			
			for(int j = 0 ; j < _rows; j++)
			{
				vector< double > & _feature = node_features[ j + node_shift];
				_feature.resize( dim_node );			
				if(IF_NODE_BIAS) _feature[dim_node-1] = 1.0;
				
				unsigned int _label;
				
				for(int k = 0 ; k < _cols ; k++)
					_feature[k] = (double) node_descs[i].at<float>(j,k);
					
				if( labels[i].at<float>(j,0) > label_threshold)
					_label = 1;
				else
					_label = 2;
					
				const RandomField::Node* _node
					=     tmp_rf.createNode( j + node_shift,
											_feature,
											_label);
											
				nodes.push_back( _node );
			
			}
			
			node_shift += _rows;
		}
		
		cout << node_shift << " nodes totally" << endl;
	}
	
	
	
	{
		// add all edge feature
		
		int node_shift = 0;	
		
		unsigned int edges_cs_idx = 0;
		
		int edge_shift = 0;
		
		for(int i = 0 ; i < (int) edge_maps.size() ; i++ )
		{
			int _rows = edge_maps[i].rows;
			
			int _cols = edge_descs[i].cols;
								
			for(int j = 0 ; j < _rows ; j++)
			{	
				vector< double > & _feature = edge_features[ j + edge_shift];
				_feature.resize( dim_edge );
				if(IF_EDGE_BIAS) _feature[dim_edge-1] = 1.0;
				
				for(int k = 0 ; k < _cols ; k++)
					_feature[k] = (double) edge_descs[i].at<float>(j,k);			
					
				int id1 = edge_maps[i].at<float>( j , 0) + node_shift;
				int id2 = edge_maps[i].at<float>( j , 1) + node_shift;
				
				list<const RandomField::Node*> _edge_nodes;
				
				_edge_nodes.push_back( nodes[ id1 ] );
				_edge_nodes.push_back( nodes[ id2 ] );
				
				tmp_rf.createClique(edges_cs_idx, _edge_nodes, _feature);
			}
			
			node_shift += node_descs[i].rows;			
			edge_shift += _rows;
		}
	}
	
	cout << "MRF node dim " << tmp_rf.getNodeDim() << endl;
	
	vector<size_t> temp_clique_size = tmp_rf.getCliqueSetDims();
	
	for(int i = 0 ; i < (int) temp_clique_size.size() ;i++)
		cout << " i: " << temp_clique_size[i] << " ";
	cout << endl;
	
	cout << "Starting to train..." << endl;	
		
	vector<const RandomField*> training_rfs(1, &tmp_rf);

	// ----------------------------------------------
	// Define learning parameters
	double nbr_iters = double(my_setting.nbr_iters);
	double step_size = my_setting.step_size;
	double truncation_param = -1.0; // indicates use Pott's model
	vector<double> robust_potts_params(tmp_rf.getNumberOfCliqueSets(), truncation_param);
	// --- Linear Model ---
	LinearRegression lin_reg(0.0001);
	FunctionalM3N m3n_model(robust_potts_params, lin_reg);
	
	if (m3n_model.train(training_rfs, nbr_iters, step_size) < 0)
	  {
		cerr << "Failed to train M3N" << endl;
		return -1;
	  }
	  cout << "Successfully trained M3N" << endl;
	  
	string dir1("learned_m3n_model");
	string dir2("resaved_m3n_model");
	string model_name = my_setting.train_video_name + "_m3n_model";
	
	if (m3n_model.saveToFile(dir1, model_name) < 0)
	{
		cerr << "Failed to save M3N" << endl;
		return -1;
	}
	FunctionalM3N loaded_m3n(dir1, model_name);
	if (loaded_m3n.saveToFile(dir2, model_name) < 0)
	{
		cerr << "Failed to load M3N" << endl;
		return -1;
	}
  
	return 0;
	
  // ----------------------------------------------------------
  // Example how to create a random field with edges and high-order cliques
  
  /*vector<double> empty_features(5); // dummy 5 dimension feature vector
  
  // Create Nodes (ids or labels do not need to be sequential)
  unsigned int node1_id = 5;
  unsigned int node1_label = 9;
  const RandomField::Node* node1 = tmp_rf.createNode(node1_id, empty_features, node1_label);
  unsigned int node2_id = 10;
  unsigned int node2_label = 50;
  const RandomField::Node* node2 = tmp_rf.createNode(node2_id, empty_features, node2_label);
  unsigned int node3_id = 500;
  unsigned int node3_label = 50;
  const RandomField::Node* node3 = tmp_rf.createNode(node3_id, empty_features, node3_label);
  unsigned int node4_id = 901;
  unsigned int node4_label = 9;
  const RandomField::Node* node4 = tmp_rf.createNode(node4_id, empty_features, node4_label);
  // Create edges (in the first clique set (cs))
  unsigned int edges_cs_idx = 0;
  list<const RandomField::Node*> edge1_nodes;
  edge1_nodes.push_back(node1);
  edge1_nodes.push_back(node2);
  tmp_rf.createClique(edges_cs_idx, edge1_nodes, empty_features);
  list<const RandomField::Node*> edge2_nodes;
  edge2_nodes.push_back(node2);
  edge2_nodes.push_back(node4);
  tmp_rf.createClique(edges_cs_idx, edge2_nodes, empty_features);
  // Create high-order cliques (in the second clique set)
  unsigned int cliques_cs_idx = 1;
  list<const RandomField::Node*> clique1_nodes;
  clique1_nodes.push_back(node1);
  clique1_nodes.push_back(node2);
  clique1_nodes.push_back(node3);
  tmp_rf.createClique(cliques_cs_idx, clique1_nodes, empty_features);
  list<const RandomField::Node*> clique2_nodes;
  clique2_nodes.push_back(node1);
  clique2_nodes.push_back(node2);
  clique2_nodes.push_back(node3);
  clique2_nodes.push_back(node4);
  tmp_rf.createClique(cliques_cs_idx, clique2_nodes, empty_features);

  // ----------------------------------------------------------
  // Load a real random field (point cloud example) to train on
  RandomField training_rf;
  if (training_rf.loadRandomFieldWGASCII("random_field/train_rf") < 0)
  {
    cerr << "Failed to load random field" << endl;
    return -1;
  }
  
  // --- Non-linear Model ---
  //OCVRTreeWrapper rtree; // default training parameters
  //FunctionalM3N m3n_model(robust_potts_params, rtree);

  // ----------------------------------------------------------
  // Train the M3N
  cout << "Starting to train..." << endl;
  if (m3n_model.train(training_rfs, nbr_iters, step_size) < 0)
  {
    cerr << "Failed to train M3N" << endl;
    return -1;
  }
  cout << "Successfully trained M3N" << endl;

  // ----------------------------------------------------------
  // Verify save and load work correctly
  string dir1("learned_m3n_model");
  string dir2("resaved_m3n_model");
  string model_name("my_m3n_model");
  if (m3n_model.saveToFile(dir1, model_name) < 0)
  {
    cerr << "Failed to save M3N" << endl;
    return -1;
  }
  FunctionalM3N loaded_m3n(dir1, model_name);
  if (loaded_m3n.saveToFile(dir2, model_name) < 0)
  {
    cerr << "Failed to load M3N" << endl;
    return -1;
  }*/

  // ----------------------------------------------------------
  // Perform inference
  /*map<unsigned int, unsigned int> inferred_labels;
  if (loaded_m3n.infer(training_rf, inferred_labels) < -0)
  {
    cerr << "Failed to perform inference" << endl;
    return -1;
  }
  const map<unsigned int, RandomField::Node*>& nodes = training_rf.getNodes();
  unsigned int first_node_id = nodes.begin()->second->getID();
  cout << "Node with id: " << first_node_id << " has predicted label: "
      << inferred_labels[first_node_id] << endl;*/

  return 0;
}
